import os
import subprocess
import sys
import argparse

def encrypt_file(input_path, output_path):
    """调用openssl直接加密单个文件"""
    KEY_SUB = "82724419649596897754226935179925"
    IV_SUB = "InitializationVe"
    
    try:
        # 转换密钥和IV为hex格式
        key_hex = KEY_SUB.encode('utf-8').hex()
        iv_hex = IV_SUB.encode('utf-8').hex()

        result = subprocess.run(
            [
                'openssl', 'enc', '-aes-256-cbc',
                '-in', input_path,
                '-out', output_path,
                '-K', key_hex,
                '-iv', iv_hex,
                '-nosalt'
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        return True
    except subprocess.CalledProcessError as e:
        raise Exception(e.stderr.strip())
    except Exception as e:
        raise e

def process_directory(input_dir, output_dir):
    """处理目录结构"""
    total = 0
    success = 0

    for root, dirs, files in os.walk(input_dir):
        # 计算相对路径
        rel_path = os.path.relpath(root, input_dir)
        output_root = os.path.join(output_dir, rel_path)
        
        # 创建输出目录
        os.makedirs(output_root, exist_ok=True)

        for filename in files:
            total += 1
            input_file = os.path.join(root, filename)
            output_filename = f"{os.path.splitext(filename)[0]}.bin"
            output_file = os.path.join(output_root, output_filename)

            try:
                encrypt_file(input_file, output_file)
                print(f"[✅] 加密成功 {input_file} -> {output_file}")
                success += 1
            except Exception as e:
                print(f"[❌] 加密失败 {input_file}")
                print(f"    错误原因: {str(e)}")

    return total, success

if __name__ == "__main__":
    # 配置命令行参数
    parser = argparse.ArgumentParser(description="文件加密工具")
    parser.add_argument('-i', '--input', required=True, help='输入目录路径')
    parser.add_argument('-o', '--output', required=True, help='输出目录路径')
    args = parser.parse_args()

    # 验证输入目录
    if not os.path.isdir(args.input):
        print(f"错误: 输入目录不存在 - {args.input}")
        sys.exit(1)

    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)

    # 执行加密
    print(f"开始加密: {args.input} => {args.output}")
    total, success = process_directory(args.input, args.output)
    
    # 输出报告
    print("\n加密完成")
    print(f"总文件数: {total}")
    print(f"成功数量: {success}")
    print(f"失败数量: {total - success}")
    sys.exit(0 if total == success else 1)